from comm_env.comm_env2 import NetworkEnv
from comm_env.scheme_env import SchemeEnvWrapper, SchemeProbEnvWrapper
from gym import register
from utils import SCHEME_LINK_PROBS, SCHEME_ENG_PROBS

register(id="Network-v0", entry_point="comm_env:NetworkEnv")
register(id="Network-v1",
         entry_point="comm_env:NetworkEnv",
         max_episode_steps=50,
         kwargs=dict(num_user=3,
                     num_channel=2,
                     link_probabilities=SCHEME_LINK_PROBS[0],
                     eng_probabilities=SCHEME_ENG_PROBS[0]))
register(id="Network-v2",
         entry_point="comm_env:NetworkEnv",
         max_episode_steps=50,
         kwargs=dict(num_user=3,
                     num_channel=2,
                     link_probabilities=SCHEME_LINK_PROBS[1],
                     eng_probabilities=SCHEME_ENG_PROBS[1]))
register(id="Network-v3",
         entry_point="comm_env:NetworkEnv",
         max_episode_steps=50,
         kwargs=dict(num_user=3,
                     num_channel=2,
                     link_probabilities=SCHEME_LINK_PROBS[2],
                     eng_probabilities=SCHEME_ENG_PROBS[2]))
register(id="Network-v4",
         entry_point="comm_env:NetworkEnv",
         max_episode_steps=50,
         kwargs=dict(num_user=3,
                     num_channel=2,
                     link_probabilities=SCHEME_LINK_PROBS[3],
                     eng_probabilities=SCHEME_ENG_PROBS[3]))
register(id="Network-v5",
         entry_point="comm_env:NetworkEnv",
         max_episode_steps=50,
         kwargs=dict(num_user=3,
                     num_channel=2,
                     link_probabilities=SCHEME_LINK_PROBS[4],
                     eng_probabilities=SCHEME_ENG_PROBS[4]))
register(id="Network-v6",
         entry_point="comm_env:NetworkEnv",
         max_episode_steps=50,
         kwargs=dict(num_user=3,
                     num_channel=2,
                     link_probabilities=SCHEME_LINK_PROBS[5],
                     eng_probabilities=SCHEME_ENG_PROBS[5]))
register(id="Network-v7",
         entry_point="comm_env:NetworkEnv",
         max_episode_steps=50,
         kwargs=dict(num_user=3,
                     num_channel=2,
                     link_probabilities=SCHEME_LINK_PROBS[6],
                     eng_probabilities=SCHEME_ENG_PROBS[6]))

__all__ = [
    "NetworkEnv",
    "SchemeEnvWrapper",
    "SchemeProbEnvWrapper",
]
